# This file executes a set of transformations 
# on a PDB file in PyMOL.  Called by getPDBImage.py.
#
# Author: emw

from glob import glob
import time

period = re.compile("\.")

for id in glob("*.pdb"):
	split_filename = period.split(id)
	pdb_id = split_filename[0]
	protein_abbreviation = split_filename[1]
	#print id
	cmd.load(pdb_id + "." + protein_abbreviation + ".pdb")
	cmd.hide("everything","all")
	cmd.show("cartoon", "all")
	#cmd.show("stick", "hetatm")
	#cmd.show("stick", "resn da+dt+dg+dc+du+hetatm")
	util.chainbow("!resn da+dt+dg+dc+du+hetatm")
	cmd.set_color("cytosine","[255,154,154]")
	cmd.set_color("guanine", "[161,161,255]")
	cmd.set_color("thymine", "[233,149,255]")
	cmd.set_color("adenine", "[120,255,120]")
	cmd.color("cytosine", "resn dc")
	cmd.color("guanine", "resn dg")
	cmd.color("thymine", "resn dt")
	cmd.color("adenine", "resn da")
	cmd.color("tv_red", "name O5'")
	cmd.color("raspberry", "name O3'")
	cmd.color("orange", "name p")
	cmd.orient(pdb_id)
	#cmd.zoom("all", "-5")
	
	# Note that background transparency doesn't show until
	# ray() is executed.
	#cmd.set("bg_rgb", [1,1,1])
	cmd.set("opaque_background", "0")
	cmd.set("show_alpha_checker","1")
	cmd.set("cartoon_transparency", "0")
	cmd.set("depth_cue", "0")
	cmd.set("ray_trace_fog", "0")
	t0 = time.time()
	cmd.ray(1200, 1000)
	print("    cmd.ray(1200, 1200) %.2f s" % (time.time() - t0))
	cmd.png("images_for_release/Protein_" + protein_abbreviation + "_PDB_" + pdb_id)

cmd.quit()